package com.nhh.v.config

import com.fasterxml.jackson.databind.ObjectMapper
import com.fasterxml.jackson.module.scala.DefaultScalaModule
import com.nhh.v.anotation.{ClientAddress, ValidateHeader}
import com.nhh.v.dto.RequiredRequestHeader
import org.springframework.beans.factory.annotation.Autowired
import org.springframework.context.annotation.Configuration
import org.springframework.core.MethodParameter
import org.springframework.http.MediaType
import org.springframework.http.codec.ServerCodecConfigurer
import org.springframework.http.codec.json.{Jackson2JsonDecoder, Jackson2JsonEncoder}
import org.springframework.web.reactive.BindingContext
import org.springframework.web.reactive.config.{EnableWebFlux, WebFluxConfigurer}
import org.springframework.web.reactive.result.method.HandlerMethodArgumentResolver
import org.springframework.web.reactive.result.method.annotation.ArgumentResolverConfigurer
import org.springframework.web.server.ServerWebExchange
import reactor.core.publisher.Mono

/**
  * Created by ywp on 2018/9/5.
  */
class HeaderArgumentResolver extends HandlerMethodArgumentResolver {

  override def resolveArgument(methodParameter: MethodParameter, bindingContext: BindingContext,
                               serverWebExchange: ServerWebExchange): Mono[AnyRef] = {
    serverWebExchange.getSession.block().getId
    val headers = serverWebExchange.getRequest.getHeaders
    val headerBuilder = RequiredRequestHeader.builder()
    if (headers.containsKey("Uid")) headerBuilder.uid(headers.get("Uid").get(0))
    if (headers.containsKey("Token")) headerBuilder.token(headers.get("Token").get(0))
    if (headers.containsKey("UserType")) headerBuilder.userType(headers.get("UserType").get(0))
    Mono.justOrEmpty(headerBuilder.build())
  }

  override def supportsParameter(methodParameter: MethodParameter): Boolean = {
    Option(methodParameter.getParameterAnnotation(classOf[ValidateHeader])).nonEmpty
  }
}

class ClientAddressArgumentResolver extends HandlerMethodArgumentResolver {

  override def resolveArgument(methodParameter: MethodParameter, bindingContext: BindingContext,
                               serverWebExchange: ServerWebExchange
                              ): Mono[AnyRef] = Mono.justOrEmpty(serverWebExchange.getRequest.getRemoteAddress.getAddress.getHostAddress)

  override def supportsParameter(methodParameter: MethodParameter): Boolean = {
    Option(methodParameter.getParameterAnnotation(classOf[ClientAddress])).nonEmpty
  }
}

@Configuration
@EnableWebFlux
class WebConfig(@Autowired val objectMapper: ObjectMapper) extends WebFluxConfigurer {

  import org.springframework.util.MimeType


  private val MIME_TYPES = Array[MimeType](
    MimeType.valueOf(MediaType.ALL_VALUE),
    MimeType.valueOf(MediaType.APPLICATION_PROBLEM_JSON_UTF8_VALUE),
    MimeType.valueOf(MediaType.APPLICATION_JSON_UTF8_VALUE)
  )

  override def configureArgumentResolvers(configurer: ArgumentResolverConfigurer): Unit = {
    configurer.addCustomResolver(new HeaderArgumentResolver)
    configurer.addCustomResolver(new ClientAddressArgumentResolver)
  }

  override def configureHttpMessageCodecs(configurer: ServerCodecConfigurer): Unit = {
    val defaultCodes = configurer.defaultCodecs()
    objectMapper.registerModule(new DefaultScalaModule)
    defaultCodes.jackson2JsonDecoder(new Jackson2JsonDecoder(objectMapper, MIME_TYPES: _*))
    defaultCodes.jackson2JsonEncoder(new Jackson2JsonEncoder(objectMapper, MIME_TYPES: _*))
  }
}
